from utils.log import logger
import json
# SUM,distinct-count,AVG
class AggRecognition:
    #初始化,data_source_path是data_dict.json的路径
    def __init__(self, model,data_source_path) -> None:
        try:
            self.model = model
            self.prompt ="假设你是一个聚合函数提取器，给你一个查询语句和一个查询中的列名，需要输出这个列名对应的聚合函数。以下是候选答案：count,avg,sum,max,min，如果该列名没有对应的聚合函数则输出none，注意不要输出解释性文本。输入查询：{}，输入列名：{}；输出：\n"
            self.agg_mode_dict={}
            with open(data_source_path, 'r', encoding='utf-8') as json_file:
                data=json.load(json_file)
                header=data[0]['header']
                for item in header:
                    if 'aggregationMode' in item.keys():
                        self.agg_mode_dict[item['id']]=item['aggregationMode']
        except Exception as err:
            logger.error("Agg recognition init error[{}]: {}".format(type(err), str(err)))
            raise
    def recognition(self,query,sel,measure):
        origin=[]
        processed=[]
        for header in sel:
            id=header['id']
            name=header['name']
            intput_prompt = self.prompt.format(query,name)
            try:
                output=self.model.predict(intput_prompt)
                origin.append({name:output})
            except Exception as err:
                logger.error("Agg recognition model output error[{}]: {}".format(type(err), str(err)))
                raise
            #默认的agg为空
            formal_output=''
            #模型没有识别到agg的情况
            try:
                if output not in ['sum','avg','count']:
                    if id in measure and id in self.agg_mode_dict.keys():
                        formal_output=self.agg_mode_dict[id]
                #模型识别到agg，要进行规范化，暂时不考虑max和min
                else:
                    if output=='sum':
                        formal_output='SUM'
                    elif output=='avg':
                        formal_output='AVG'
                    elif output=='count':
                        formal_output='distinct-count'
            except Exception as err:
                logger.error("Agg recognition standardization error[{}]: {}".format(type(err), str(err)))
                raise
            processed.append(formal_output)
        return {'origin':origin,'processed':{'agg':processed}}
            



    